/*
<%
cfg['compiler_args'] = ['-std=c++11', '-undefined dynamic_lookup']
%>
<%
setup_pybind11(cfg)
%>
*/
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <iostream>
#include <random>
#include <algorithm>
#include <time.h>

typedef unsigned int ui;

using namespace std;
namespace py = pybind11;

int randint_(int end)
{
    return rand() % end;
}

py::array_t<int> sample_negative(int user_num, int item_num, int train_num, std::vector<std::vector<int>> allPos, int neg_num)
{
    int perUserNum = (train_num / user_num);
    int row = neg_num + 2;
    py::array_t<int> S_array = py::array_t<int>({user_num * perUserNum, row});
    py::buffer_info buf_S = S_array.request();
    int *ptr = (int *)buf_S.ptr;

    for (int user = 0; user < user_num; user++)
    {
        std::vector<int> pos_item = allPos[user];

        for (int pair_i = 0; pair_i < perUserNum; pair_i++)
        {
            int negitem = 0;
            ptr[(user * perUserNum + pair_i) * row] = user;
            ptr[(user * perUserNum + pair_i) * row + 1] = pos_item[randint_(pos_item.size())];
            for (int index = 2; index < neg_num + 2; index++)
            {
                do
                {
                    negitem = randint_(item_num);     // 核心，随机负采样
                } while (
                    find(pos_item.begin(), pos_item.end(), negitem) != pos_item.end());
                ptr[(user * perUserNum + pair_i) * row + index] = negitem;
            }
        }
    }
    return S_array;
}

//py::array_t<int> sample_negative_ByUser(std::vector<int> users, int item_num, std::vector<std::vector<int>> allPos, int neg_ratio)
//{
//    int row = neg_ratio + 2;
//    int col = users.size();
//    py::array_t<int> S_array = py::array_t<int>({col, row});
//    py::buffer_info buf_S = S_array.request();
//    int *ptr = (int *)buf_S.ptr;
//
//    for (int user_i = 0; user_i < users.size(); user_i++)
//    {
//        int user = users[user_i];
//        std::vector<int> pos_item = allPos[user];
//        int negitem = 0;
//
//        ptr[user_i * row] = user;
//        ptr[user_i * row + 1] = pos_item[randint_(pos_item.size())];
//
//        for (int neg_i = 2; neg_i < row; neg_i++)
//        {
//            do
//            {
//                negitem = randint_(item_num);
//            } while (
//                find(pos_item.begin(), pos_item.end(), negitem) != pos_item.end());
//            ptr[user_i * row + neg_i] = negitem;
//        }
//    }
//    return S_array;
//}


// 按用户笛卡尔积采样，负样本率为1
py::array_t<int> sample_negative_ByUser(int user_num, int item_num, int train_num,
                                        std::vector<std::vector<int>> allPos)
{
    // int row = neg_ratio + 2; // user, pos_item, neg_item
    py::array_t<int> S_array = py::array_t<int>({train_num, 3});   // train_num 表示每个用户正负样本集合笛卡尔积，对所有用户求和后的数量
    py::buffer_info buf_S = S_array.request();
    int *ptr = (int *)buf_S.ptr;
    int total = 0;

    for (int user_i = 0; user_i < user_num; user_i++)
    {
        std::vector<int> pos_item = allPos[user_i];
        unsigned int pos_item_size = pos_item.size();
        if(pos_item_size == 0) continue;
        int negitem = 0;
        int cur_index = 0;
		std::vector<int> neg_item(pos_item_size, 0);

        // 负采样|P|个负样本
		for (unsigned int neg_i = 0; neg_i < pos_item_size; neg_i++)
			{
				do
				{
					negitem = randint_(item_num);
				} while (
					find(pos_item.begin(), pos_item.end(), negitem) != pos_item.end());
				neg_item[neg_i] = negitem;
			}

        // 把 pos_item.size()的平方个数填入到buf_S中去。
		for (unsigned int pos_i = 0; pos_i < pos_item_size; pos_i++){
			for (unsigned int neg_i = 0; neg_i < pos_item_size; neg_i++){
                cur_index = total + (pos_i*pos_item_size + neg_i)*3;
				ptr[cur_index] = user_i;
				ptr[cur_index + 1] = pos_item[pos_i];
				ptr[cur_index + 2] = neg_item[neg_i];

			}
		}
        total += pos_item_size*pos_item_size*3;
//        std::cout << total << "    "  << pos_item_size << endl;
    }
    return S_array;
}


// 按用户正样本1:1采样，构成 user-item-label 样本风格的数据集
py::array_t<int> pointwise_negtive_sample(int user_num, int item_num, int train_num,
                                        std::vector<std::vector<int>> allPos)
{
    // int row = neg_ratio + 2; // user, pos_item, neg_item
    py::array_t<int> S_array = py::array_t<int>({2*train_num, 3});   // train_num 表示所有用户正样本总数
    py::buffer_info buf_S = S_array.request();
    int *ptr = (int *)buf_S.ptr;
    int total = 0;

    for (int user_i = 0; user_i < user_num; user_i++)
    {
        std::vector<int> pos_item = allPos[user_i];
        unsigned int pos_item_size = pos_item.size();
        if(pos_item_size == 0) continue;
        int neg_item = 0;
        int cur_index = 0;

        // 把 pos_item.size()的平方个数填入到buf_S中去。
		for (unsigned int pos_i = 0; pos_i < pos_item_size; pos_i++){
            do{
            neg_item = randint_(item_num);
            } while (find(pos_item.begin(), pos_item.end(), neg_item) != pos_item.end());

            cur_index = total + pos_i*2*3;
            ptr[cur_index] = user_i;
            ptr[cur_index + 1] = pos_item[pos_i];
            ptr[cur_index + 2] = 1;
            ptr[cur_index + 3] = user_i;
            ptr[cur_index + 4] = neg_item;
            ptr[cur_index + 5] = 0;
		}
        total += pos_item_size*2*3;
    }
    return S_array;
}


void set_seed(unsigned int seed)
{
    srand(seed);
}

using namespace py::literals;

PYBIND11_MODULE(sampling, m)
{
    srand(time(0));
    // srand(2020);
    m.doc() = "example plugin";
    m.def("randint", &randint_, "generate int between [0 end]", "end"_a);
    m.def("seed", &set_seed, "set random seed", "seed"_a);
    m.def("sample_negative", &sample_negative, "sampling negatives for all",
          "user_num"_a, "item_num"_a, "train_num"_a, "allPos"_a, "neg_num"_a);
    m.def("sample_negative_ByUser", &sample_negative_ByUser, "sampling negatives for given users",
          "user_num"_a, "item_num"_a, "train_num"_a, "allPos"_a);
    m.def("pointwise_negtive_sample", &pointwise_negtive_sample, "sampling negatives for given users",
          "user_num"_a, "item_num"_a, "train_num"_a, "allPos"_a);
}
